# 写一个向gpt-4o-audio-preview发送音频的程序
# 音频使用 .input/hello_en.wav

import os
import base64
import numpy as np
import soundfile as sf
import time
from openai import OpenAI
from dotenv import load_dotenv, find_dotenv


def encode_audio_file(file_path):
    """将音频文件编码为 base64 字符串"""
    with open(file_path, "rb") as audio_file:
        return base64.b64encode(audio_file.read()).decode("utf-8")


def load_config():
    """加载环境变量配置"""
    load_dotenv(find_dotenv(), override=True)

    api_key = os.getenv("CHATFIRE_OPENAI_API_KEY")
    base_url = os.getenv("CHATFIRE_OPENAI_BASE_URL")

    if not api_key or not base_url:
        raise ValueError(
            "请确保设置了 CHATFIRE_OPENAI_API_KEY 和 CHATFIRE_OPENAI_BASE_URL 环境变量"
        )

    return api_key, base_url


def validate_audio_file(file_path):
    """验证音频文件是否存在"""
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"音频文件 {file_path} 不存在")
    return True


def create_openai_client(api_key, base_url):
    """创建 OpenAI 客户端"""
    return OpenAI(api_key=api_key, base_url=base_url)


def prepare_audio_data(file_path):
    """准备音频数据"""
    print(f"正在处理音频文件: {file_path}")

    try:
        audio_data = encode_audio_file(file_path)
        print("音频文件编码成功")
        return audio_data
    except Exception as e:
        raise RuntimeError(f"音频文件编码失败: {e}")


def create_chat_completion(
    client, audio_data, prompt=None, model="gpt-4o-audio-preview"
):
    """创建聊天完成请求"""
    prompt = (
        prompt
        or "As customer support, detect the spoken language first and reply kindly in that same language."
    )
    return client.chat.completions.create(
        model=model,
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "input_audio",
                        "input_audio": {
                            "data": audio_data,
                            "format": "wav",
                        },
                    },
                    {
                        "type": "text",
                        "text": prompt,
                    },
                ],
            },
        ],
        modalities=["text", "audio"],
        audio={"voice": "shimmer", "format": "pcm16"},
        stream=True,
        stream_options={"include_usage": True},
    )


def process_stream_response(completion):
    """处理流式响应"""
    print("\n开始接收响应...")
    print("文本响应: ", end="")

    max_tokens = 200
    outputlen = 0
    audio_string = ""
    start_audio_time = None

    for chunk in completion:
        if chunk.choices:
            # 处理音频数据
            if (
                hasattr(chunk.choices[0].delta, "audio")
                and chunk.choices[0].delta.audio
            ):
                try:
                    # 收集音频数据
                    if "data" in chunk.choices[0].delta.audio:
                        print("🎵", end="")
                        audio_string += chunk.choices[0].delta.audio["data"]
                        if start_audio_time is None:
                            start_audio_time = time.time()

                    # 显示音频转录文本
                    if "transcript" in chunk.choices[0].delta.audio:
                        print(chunk.choices[0].delta.audio["transcript"], end="")
                        outputlen += len(chunk.choices[0].delta.audio["transcript"])
                        if outputlen > max_tokens:
                            outputlen = 0
                            print("\n")
                except Exception as e:
                    print(f"\n处理音频数据时出错: {e}")

            # 处理文本响应
            if (
                hasattr(chunk.choices[0].delta, "content")
                and chunk.choices[0].delta.content
            ):
                print(chunk.choices[0].delta.content, end="")
        else:
            # 显示使用统计
            if hasattr(chunk, "usage") and chunk.usage:
                print(f"\n\n使用统计: {chunk.usage}")

    print("\n")
    return audio_string, start_audio_time


def save_audio_response(audio_string, output_filename="response_hello_en.wav"):
    """保存音频响应到文件"""
    if not audio_string:
        print("没有音频响应数据")
        return None

    try:
        # 确保输出目录存在
        os.makedirs(".output", exist_ok=True)

        # 解码并保存音频 (PCM16 格式)
        wav_bytes = base64.b64decode(audio_string)
        # PCM16 格式：16位有符号整数，小端序
        audio_np = np.frombuffer(wav_bytes, dtype=np.int16)
        output_path = f".output/{output_filename}"
        # PCM16 格式通常使用 24000Hz 采样率
        sf.write(output_path, audio_np, samplerate=24000)
        print(f"音频响应已保存到: {output_path}")
        return output_path
    except Exception as e:
        print(f"保存音频响应时出错: {e}")
        return None


def print_timing_stats(start_time, start_audio_time):
    """打印时间统计信息"""
    total_time = time.time() - start_time
    print(f"总耗时: {total_time:.2f} 秒")

    if start_audio_time:
        audio_time = time.time() - start_audio_time
        print(f"音频处理耗时: {audio_time:.2f} 秒")


def process_audio_request(
    audio_file_path=".input/hello_en.wav", output_filename="response_hello_en.wav"
):
    """处理音频请求的主要流程"""
    try:
        # 1. 加载配置
        api_key, base_url = load_config()
        print(f"使用 API: {base_url}")

        # 2. 验证音频文件
        validate_audio_file(audio_file_path)

        # 3. 创建客户端
        client = create_openai_client(api_key, base_url)

        # 4. 准备音频数据
        audio_data = prepare_audio_data(audio_file_path)

        # 5. 记录开始时间
        start_time = time.time()

        # 6. 创建聊天完成请求
        completion = create_chat_completion(
            client,
            audio_data,
            prompt="这段音频在说什么",
            # prompt="As customer support, detect the spoken language first and reply kindly in that same language.",
        )

        # 7. 处理流式响应
        audio_string, start_audio_time = process_stream_response(completion)

        # 8. 保存音频响应
        save_audio_response(audio_string, output_filename)

        # 9. 显示时间统计
        print_timing_stats(start_time, start_audio_time)

    except ValueError as e:
        print(f"配置错误：{e}")
    except FileNotFoundError as e:
        print(f"文件错误：{e}")
    except RuntimeError as e:
        print(f"处理错误：{e}")
    except Exception as e:
        print(f"请求失败: {e}")


def main():
    """主函数"""
    process_audio_request(
        audio_file_path=".input/cherry.wav", output_filename="response_cherry.wav"
    )


if __name__ == "__main__":
    main()
